# @Time : 2021/8/5 17:23
# @Author : Li Kunlun
# @Description : 使用重复元素的网络（VGG）

import utils as d2l
from mxnet import gluon, init, nd
from mxnet.gluon import nn


# 1、VGG块
# 连续使用数个相同的填充为1、窗口形状为 3×3 的卷积层后接上一个步幅为2、窗口形状为 2×2 的最大池化层。
def vgg_block(num_convs, num_channels):
    blk = nn.Sequential()
    for _ in range(num_convs):
        blk.add(nn.Conv2D(num_channels, kernel_size=3, padding=1, activation='relu'))
    blk.add(nn.MaxPool2D(pool_size=2, strides=2))
    return blk


# 2、VGG网络
# 该变量指定了每个VGG块里卷积层个数和输出通道数
conv_arch = ((1, 64), (1, 128), (2, 256), (2, 512), (2, 512))


# VGG-11
def vgg(conv_arch):
    net = nn.Sequential()
    # 卷积层部分
    for (num_convs, num_channels) in conv_arch:
        net.add(vgg_block(num_convs, num_channels))
    # 全连接层部分
    net.add(nn.Dense(4096, activation='relu'), nn.Dropout(0.5),
            nn.Dense(4096, activation='relu'), nn.Dropout(0.5),
            nn.Dense(10))
    return net


net = vgg(conv_arch)

net.initialize()
X = nd.random.uniform(shape=(1, 1, 224, 224))
for blk in net:
    X = blk(X)
    """
    sequential1 output shape:	 (1, 64, 112, 112)
    sequential2 output shape:	 (1, 128, 56, 56)
    sequential3 output shape:	 (1, 256, 28, 28)
    sequential4 output shape:	 (1, 512, 14, 14)
    sequential5 output shape:	 (1, 512, 7, 7)
    dense0 output shape:	 (1, 4096)
    dropout0 output shape:	 (1, 4096)
    dense1 output shape:	 (1, 4096)
    dropout1 output shape:	 (1, 4096)
    dense2 output shape:	 (1, 10)
    """
    print(blk.name, 'output shape:\t', X.shape)

# 3、训练模型
# 出于测试的目的我们构造一个通道数更小，或者说更窄的网络在Fashion-MNIST数据集上进行训练。
ratio = 4
small_conv_arch = [(pair[0], pair[1] // ratio) for pair in conv_arch]
net = vgg(small_conv_arch)

# 使用了稍大些的学习率，模型训练过程与上一节的AlexNet中的类似。
lr, num_epochs, batch_size, ctx = 0.05, 5, 128, d2l.try_gpu()
net.initialize(ctx=ctx, init=init.Xavier())
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': lr})
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size, resize=224)
d2l.train_ch5(net, train_iter, test_iter, batch_size, trainer, ctx,num_epochs)
